from gurobipy import *
import numpy as np
from traffic_engineer import *
from traffic_engineering_class import *

# given a topology, and the traffic matrix, will compute the best case
class IdealTrafficEngineer(TrafficEngineer):
	def __init__(self, aurora_network, all_interblock_paths):
		TrafficEngineer.__init__(self, aurora_network, all_interblock_paths)
		return

	def get_filename_param(self):
		return "perfect"

	# returns a string description of this class
	def get_string(self):
		return "Perfect TE (Oracle Prediction)"

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.IDEAL

	def compute_path_weights(self, logical_topology_adj_matrix, traffic_matrix):
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("Routing minimize MLU for traffic matrix")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="mlu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * nblocks
		for i in range(nblocks):
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					link_capacity_constraints[i][j] = LinExpr()
					weights_sum_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=0., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weights_sum_constraint += routing_weight_vars[path]
					model.addConstr(lhs=weights_sum_constraint, sense=GRB.EQUAL, rhs=1.)

		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					for path in self.all_interblock_paths[i][j]:
						src = path[0]
						dst = path[-1]
						curr_node = path[0]
						for path_hop in range(1, len(path), 1):
							next_node = path[path_hop]
							link_capacity_constraints[curr_node][next_node] += (routing_weight_vars[path] * traffic_matrix[i][j])
							curr_node = next_node

		## stage 2 : add the link utilization constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=link_capacity_constraints[i][j], sense=GRB.LESS_EQUAL, rhs=mlu * capacity * logical_topology_adj_matrix[i][j])
		# stage 3: set the objective
		model.setObjective(mlu, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for path in routing_weight_vars.keys():
				routing_weights[path] = routing_weight_vars[path].x
			return routing_weights
		except GurobiError as e1:
			print ("Error code " + str(e1. errno ) + ": " + str(e1))
			return None
		except AttributeError as e:
			print ("Ideal Traffic Engineering - Encountered an attribute error " + ": " + str(e))
			return None

	# minimizes MLU subject to the fact that hop count equals to the argument
	def compute_path_weights_fix_avehopcount(self, logical_topology_adj_matrix, traffic_matrix, avehopcount):
		assert(avehopcount >= 1 and avehopcount <= 2)
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("Routing minimize MLU for traffic matrix")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="mlu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * nblocks
		for i in range(nblocks):
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					link_capacity_constraints[i][j] = LinExpr()
					weights_sum_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0., ub=GRB.INFINITY, obj=0., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weights_sum_constraint += routing_weight_vars[path]
					model.addConstr(lhs=traffic_matrix[i][j], sense=GRB.EQUAL, rhs=weights_sum_constraint)

		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					for path in self.all_interblock_paths[i][j]:
						src = path[0]
						dst = path[-1]
						curr_node = path[0]
						for path_hop in range(1, len(path), 1):
							next_node = path[path_hop]
							link_capacity_constraints[curr_node][next_node] += routing_weight_vars[path]
							curr_node = next_node

		## stage 2 : add the link utilization constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=link_capacity_constraints[i][j], sense=GRB.LESS_EQUAL, rhs=mlu * capacity * logical_topology_adj_matrix[i][j])

		## stage 2.5 : fix the hop count
		hopcount_constraint = LinExpr()
		total_traffic = 0.
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j and traffic_matrix[i][j] > 0:
					total_traffic += traffic_matrix[i][j]
					for path in self.all_interblock_paths[i][j]:
						hc = 1.
						if len(path) > 2:
							hc = 2.
						hopcount_constraint += (hc * routing_weight_vars[path])
		model.addConstr(lhs=hopcount_constraint, sense=GRB.EQUAL, rhs=avehopcount * total_traffic)
		#print("total traffic is : {} avehopcount is: {}".format(total_traffic, avehopcount))
		#model.addConstr(lhs=0 * avehopcount * total_traffic, sense=GRB.LESS_EQUAL, rhs=hopcount_constraint)
		#model.addConstr(lhs=hopcount_constraint, sense=GRB.LESS_EQUAL, rhs=100. * avehopcount * total_traffic)

		# stage 3: set the objective
		model.setObjective(mlu, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						if traffic_matrix[i][j] > 0:
							traffic_sum = 0.
							for path in self.all_interblock_paths[i][j]:
								traffic_sum += routing_weight_vars[path].x
								routing_weights[path] = routing_weight_vars[path].x
							for path in self.all_interblock_paths[i][j]:
								routing_weights[path] /= traffic_sum
						else:
							# ecmp otherwise
							ratio = 1./len(self.all_interblock_paths[i][j])
							for path in self.all_interblock_paths[i][j]:
								routing_weights[path] = ratio
			return routing_weights
		except GurobiError as e1:
			print ("Error code " + str(e1. errno ) + ": " + str(e1))
			return None
		except AttributeError as e:
			print ("Ideal Traffic Engineering - same hop count Encountered an attribute error " + ": " + str(e))
			return None